//! Virtual memory space management in Zig
//! SPDX-License-Identifier: MPL-2.0

const std = @import("std");
const lib = @import("lib.zig");
const frame = @import("frame.zig");
const page_table = @import("page_table.zig");

const Vaddr = lib.Vaddr;
const Paddr = lib.Paddr;
const PageFlags = lib.PageFlags;
const MemoryError = lib.MemoryError;
const PAGE_SIZE = lib.PAGE_SIZE;

// Virtual memory area
pub const VmArea = struct {
    start: Vaddr,
    end: Vaddr,
    flags: PageFlags,
    next: ?*VmArea,

    const Self = @This();

    pub fn init(start: Vaddr, end: Vaddr, flags: PageFlags) Self {
        return Self{
            .start = start,
            .end = end,
            .flags = flags,
            .next = null,
        };
    }

    pub fn size(self: Self) usize {
        return self.end - self.start;
    }

    pub fn contains(self: Self, addr: Vaddr) bool {
        return addr >= self.start and addr < self.end;
    }

    pub fn overlaps(self: Self, other_start: Vaddr, other_end: Vaddr) bool {
        return !(self.end <= other_start or self.start >= other_end);
    }
};

// Virtual memory space
pub const VmSpace = struct {
    page_table: page_table.PageTable,
    vm_areas: ?*VmArea,
    mutex: std.Thread.Mutex,
    allocator: std.mem.Allocator,

    const Self = @This();

    pub fn init(allocator: std.mem.Allocator) Self {
        return Self{
            .page_table = page_table.PageTable.init(),
            .vm_areas = null,
            .mutex = std.Thread.Mutex{},
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *Self) void {
        self.mutex.lock();
        defer self.mutex.unlock();

        // Free all VM areas
        var current = self.vm_areas;
        while (current) |area| {
            const next = area.next;
            self.allocator.destroy(area);
            current = next;
        }
    }

    pub fn map(self: *Self, vaddr: Vaddr, paddr: Paddr, size: usize, flags: PageFlags, frame_allocator: frame.FrameAllocator) !void {
        if (!lib.isPageAligned(vaddr) or !lib.isPageAligned(paddr) or !lib.isPageAligned(size)) {
            return MemoryError.InvalidAlignment;
        }

        self.mutex.lock();
        defer self.mutex.unlock();

        // Check for overlaps with existing areas
        if (self.findOverlappingArea(vaddr, vaddr + size)) |_| {
            return MemoryError.AlreadyMapped;
        }

        // Create VM area
        const vm_area = try self.allocator.create(VmArea);
        vm_area.* = VmArea.init(vaddr, vaddr + size, flags);
        self.insertVmArea(vm_area);

        // Map pages
        const walker = page_table.PageTableWalker.init(&self.page_table);
        var current_vaddr = vaddr;
        var current_paddr = paddr;
        const end_addr = vaddr + size;

        while (current_vaddr < end_addr) {
            try walker.map(current_vaddr, current_paddr, flags, frame_allocator, self.allocator);
            current_vaddr += PAGE_SIZE;
            current_paddr += PAGE_SIZE;
        }
    }

    pub fn unmap(self: *Self, vaddr: Vaddr, size: usize) !void {
        if (!lib.isPageAligned(vaddr) or !lib.isPageAligned(size)) {
            return MemoryError.InvalidAlignment;
        }

        self.mutex.lock();
        defer self.mutex.unlock();

        const end_addr = vaddr + size;
        
        // Find and remove VM area
        if (self.findAndRemoveArea(vaddr, end_addr)) |area| {
            self.allocator.destroy(area);
        } else {
            return MemoryError.NotMapped;
        }

        // Unmap pages
        const walker = page_table.PageTableWalker.init(&self.page_table);
        var current_vaddr = vaddr;

        while (current_vaddr < end_addr) {
            try walker.unmap(current_vaddr);
            current_vaddr += PAGE_SIZE;
        }
    }

    pub fn translate(self: *Self, vaddr: Vaddr) ?Paddr {
        self.mutex.lock();
        defer self.mutex.unlock();

        const walker = page_table.PageTableWalker.init(&self.page_table);
        return walker.translate(vaddr);
    }

    pub fn findVmArea(self: *Self, vaddr: Vaddr) ?*VmArea {
        self.mutex.lock();
        defer self.mutex.unlock();

        var current = self.vm_areas;
        while (current) |area| {
            if (area.contains(vaddr)) {
                return area;
            }
            current = area.next;
        }
        return null;
    }

    fn findOverlappingArea(self: *Self, start: Vaddr, end: Vaddr) ?*VmArea {
        var current = self.vm_areas;
        while (current) |area| {
            if (area.overlaps(start, end)) {
                return area;
            }
            current = area.next;
        }
        return null;
    }

    fn insertVmArea(self: *Self, new_area: *VmArea) void {
        if (self.vm_areas == null or new_area.start < self.vm_areas.?.start) {
            new_area.next = self.vm_areas;
            self.vm_areas = new_area;
            return;
        }

        var current = self.vm_areas;
        while (current) |area| {
            if (area.next == null or new_area.start < area.next.?.start) {
                new_area.next = area.next;
                area.next = new_area;
                return;
            }
            current = area.next;
        }
    }

    fn findAndRemoveArea(self: *Self, start: Vaddr, end: Vaddr) ?*VmArea {
        if (self.vm_areas) |first_area| {
            if (first_area.start == start and first_area.end == end) {
                self.vm_areas = first_area.next;
                return first_area;
            }
        }

        var current = self.vm_areas;
        while (current) |area| {
            if (area.next) |next_area| {
                if (next_area.start == start and next_area.end == end) {
                    area.next = next_area.next;
                    return next_area;
                }
            }
            current = area.next;
        }
        return null;
    }
};

// Tests
test "vm area creation" {
    const testing = std.testing;
    const flags = PageFlags{ .present = true, .writable = true };
    const area = VmArea.init(0x1000, 0x5000, flags);
    
    try testing.expectEqual(@as(Vaddr, 0x1000), area.start);
    try testing.expectEqual(@as(Vaddr, 0x5000), area.end);
    try testing.expectEqual(@as(usize, 0x4000), area.size());
    try testing.expect(area.contains(0x3000));
    try testing.expect(!area.contains(0x6000));
}

test "vm area overlap detection" {
    const testing = std.testing;
    const flags = PageFlags{ .present = true };
    const area = VmArea.init(0x1000, 0x5000, flags);
    
    try testing.expect(area.overlaps(0x0000, 0x2000)); // Partial overlap at start
    try testing.expect(area.overlaps(0x3000, 0x6000)); // Partial overlap at end
    try testing.expect(area.overlaps(0x2000, 0x3000)); // Completely inside
    try testing.expect(area.overlaps(0x0000, 0x6000)); // Completely contains
    try testing.expect(!area.overlaps(0x6000, 0x7000)); // No overlap
}